from backend.api import (
    get_factor_list,
    get_factor_rank,
    get_factor_values,
    get_symbol_list,
)
from page_utils import launch_task_widget
import pandas as pd
import plotly.express as px
import streamlit as st


if __name__ == "__main__":
    st.title("因子任务")
    with st.form("factor_calculate"):
        symbols = st.multiselect("Symbol", get_symbol_list())
        factors = st.multiselect("Factor", get_factor_list())

        start = st.date_input("Start Date")
        end = st.date_input("End Date")
        timeframe = st.selectbox("Timeframe", ["1m", "5m", "15m", "1h", "1d"])

        agg_method = st.selectbox("Aggregate Method", ["因子聚合", "品种聚合"], index=0)
        sbm = st.form_submit_button("绘制因子曲线图")

    if sbm:
        factor_values = get_factor_values(symbols, factors, start, end, timeframe)
        # 根据factor_values的值，使用streamlit在前端页面中绘制因子曲线图

        if not factor_values.empty:
            if agg_method == "品种聚合":
                for factor in factors:
                    df = pd.DataFrame(
                        {symbol: factor_values[symbol][factor] for symbol in symbols}
                    )
                    st.subheader(f"{factor} Factor")
                    fig = px.line(
                        df,
                        title=f"{factor} Factor Curve",
                        labels={
                            "index": "Time",
                            "value": "Value",
                            "variable": "Symbol",
                        },
                    )
                    st.plotly_chart(fig)
            else:
                for symbol in symbols:
                    df = pd.DataFrame(
                        {symbol: factor_values[symbol][factor] for factor in factors}
                    )
                    st.subheader(f"{symbol} Symbol")
                    fig = px.line(
                        df,
                        title=f"{symbol} Symbol Curve",
                        labels={
                            "index": "Time",
                            "value": "Value",
                            "variable": "Factor",
                        },
                    )
                    st.plotly_chart(fig)
    if st.button("显示因子排名"):
        factor_rank = get_factor_rank(symbols, factors, start, end)
        st.dataframe(factor_rank)

    launch_task_widget("factor_calculate", symbols, factors, start, end, timeframe)
